package main

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"

	"entgo.io/ent/entc"
	"entgo.io/ent/entc/gen"
)

var _ entc.Extension = (*RepoExtension)(nil)

// RepoExtension 是一个简单的 entc 扩展示例：
// 它在正常代码生成完成后，写入一个 `repositories_gen.go` 文件，
// 为每个实体生成一个极简的 Repository 接口（示例用途）。
type RepoExtension struct {
	entc.DefaultExtension
	// Module is program package name.
	Module string
	// OutDir is the base directory where the extension will write files.
	// For example: "pkg/repository". The extension will create one subdir
	// per schema under this base dir (e.g. pkg/repository/user).
	OutDir string
	// FileMode controls the permissions for created files (defaults to 0644).
	FileMode os.FileMode
	// Template is the parsed gen.Template used to render repository files.
	Template *gen.Template
}

// Hooks 返回我们注册的 Hook：在主生成流程完成后写文件。
func (r RepoExtension) Hooks() []gen.Hook {
	hook := func(next gen.Generator) gen.Generator {
		return gen.GenerateFunc(func(g *gen.Graph) error {
			// 首先执行默认生成器
			if err := next.Generate(g); err != nil {
				return err
			}
			// 为每个 schema 生成独立目录并写文件到 r.OutDir/<snake_name>/
			base := r.OutDir
			if base == "" {
				// 如果未配置，从 gen.Config.Target 推断（保持兼容旧行为）
				base = g.Config.Target
			}
			if base == "" {
				// 最后回退到默认目录
				base = "ent"
			}

			for _, t := range g.Nodes {
				// 目录名使用驼峰转蛇形（例如 User -> user）
				dirName := ToSnakeCase(t.Name)
				outDir := filepath.Join(base, dirName)
				if err := os.MkdirAll(outDir, 0o755); err != nil {
					return err
				}

				var buf bytes.Buffer
				if r.Template != nil {
					// Execute named template with a single-node graph data.
					data := struct {
						Module string
						Nodes  []*gen.Type
						Config *gen.Config
					}{
						Module: r.Module,
						Nodes:  []*gen.Type{t},
						Config: g.Config,
					}
					// The template file defines a template named "repository/dao".
					if err := r.Template.ExecuteTemplate(&buf, "repository/dao", data); err != nil {
						return err
					}
				}

				fname := filepath.Join(outDir, fmt.Sprintf("%s_gen.go", ToSnakeCase(t.Name)))
				mode := r.FileMode
				if mode == 0 {
					mode = 0o644
				}
				if err := os.WriteFile(fname, buf.Bytes(), mode); err != nil {
					return err
				}
				fmt.Println("RepoExtension: wrote", fname)
			}
			return nil
		})
	}
	return []gen.Hook{hook}
}

// NewRepoExtension 返回默认扩展实例，传入 outDir 指定基准输出目录（例如 "pkg/repository"），
// template 用于渲染每个实体的文件（如果为 nil 将回退到内建的文本渲染）。
func NewRepoExtension(module, outDir string, template *gen.Template) entc.Extension {
	return RepoExtension{Module: module, OutDir: outDir, Template: template}
}
